# -*- coding: utf-8 -*-
"""
Created on Thu Apr  7 10:44:10 2022

@author: 11325
"""
import torch
import numpy as np
import pandas as pd
import torch.autograd as autograd
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import dataloader 
import os
import tqdm


class Evaluation():
    def __init__(self,model,test_tuple):
        self.model=model
        self.test_tuple=test_tuple
        self.meanrank=0
        self.hit10=0
        
        self.relation_meanrank=0
        self.relation_hit10=0
        
        
    def transE_testdata(self,test_pair):
        test_pair=torch.tensor(test_pair).cuda()
        head_embedding=torch.unsqueeze(self.model.enlity(test_pair[0]), 0)
        tail_embedding=torch.unsqueeze(self.model.enlity(test_pair[1]), 0)
        relation_embedding=torch.unsqueeze(self.model.relation(test_pair[2]), 0)
        return head_embedding,tail_embedding,relation_embedding
    
    def transH_or_R_testdata(self,test_pair):
        test_pair=torch.tensor(test_pair).cuda()
        head_embedding=torch.unsqueeze(self.model.enlity(test_pair[0]), 0)
        tail_embedding=torch.unsqueeze(self.model.enlity(test_pair[1]), 0)
        relation_embedding=torch.unsqueeze(self.model.relation(test_pair[2]), 0)
        relation_wr=torch.unsqueeze(self.model.relation_Wr(test_pair[:,2]),0)
        return head_embedding,tail_embedding,relation_embedding,relation_wr
    
        
    def pair_rank_transE(self,pair):
        hit10_head=0
        hit10_tail=0
        
        test_head,test_tail,test_relation=self.transE_testdata(pair)
        #计算输入的测试元组的score
        #print(self.enlity_len)
        
        loss_test_pair=self.model.score(test_head,test_tail,test_relation)        
        #测试元组增维
        
        test_head=test_head.repeat(self.model.enlity_len,1)
        test_tail=test_tail.repeat(self.model.enlity_len,1)
        test_relation=test_relation.repeat(self.model.enlity_len,1)
        
        
        #计算对头节点的预测分数
        loss_head=self.model.score(self.model.enlity.weight.data,test_tail,test_relation)
        loss_tail=self.model.score(test_head,self.model.enlity.weight.data,test_relation)
        #loss_relation=torch.norm(test_head_embeeding+self.relation_embedding.weight.data-test_tail_embeeding,2,1).view(-1,1)
        c= loss_test_pair-loss_head
        #计算mean_rank
        head_rank=torch.nonzero(nn.functional.relu(loss_test_pair-loss_head))[:,0]
        tail_rank=torch.nonzero(nn.functional.relu(loss_test_pair-loss_tail))[:,0]
        #relation_rank=torch.nonzero(nn.functional.relu(loss_test_pair-loss_relation))[:,0]
        
        num_false_head=head_rank.size()[0]#比测试元组分数小的头元组个数
        num_false_tail=tail_rank.size()[0]#比测试元组分数小的尾元组个数
        #num_false_relation=relation_rank.size()[0]#比测试元组分数小的关系元组个数
        
        mean_rank_head=num_false_head
        mean_rank_tail=num_false_tail
        #mean_rank_relation=num_false_relation+1
        
        #计算hit10指标
        if (mean_rank_head<=10):
            hit10_head=1
        if (mean_rank_tail<=10):
            hit10_tail=1
            '''
        if (mean_rank_relation<10):
            hit10_relation=1
            '''
        return mean_rank_head,mean_rank_tail,hit10_head,hit10_tail
    
    def pair_rank_transH_or_R(self,pair):
        hit10_head=0
        hit10_tail=0
        
        test_head,test_tail,test_relation,test_relation_wr=self.transH_or_R_testdata(pair)
        #计算输入的测试元组的score
        #print(self.enlity_len)
        
        loss_test_pair=self.model.score(test_head,test_tail,test_relation,test_relation_wr)        
        #测试元组增维
        
        test_head=test_head.repeat(self.model.enlity_len,1)
        test_tail=test_tail.repeat(self.model.enlity_len,1)
        test_relation=test_relation.repeat(self.model.enlity_len,1)
        test_relation_wr=test_relation_wr.repeat(self.model.enlity_len,1)
        
        
        #计算对头节点的预测分数
        loss_head=self.model.score(self.model.enlity.weight.data,test_tail,test_relation,test_relation_wr)
        loss_tail=self.model.score(test_head,self.model.enlity.weight.data,test_relation,test_relation_wr)
        #loss_relation=torch.norm(test_head_embeeding+self.relation_embedding.weight.data-test_tail_embeeding,2,1).view(-1,1)
        c= loss_test_pair-loss_head
        #计算mean_rank
        head_rank=torch.nonzero(nn.functional.relu(loss_test_pair-loss_head))[:,0]
        tail_rank=torch.nonzero(nn.functional.relu(loss_test_pair-loss_tail))[:,0]
        #relation_rank=torch.nonzero(nn.functional.relu(loss_test_pair-loss_relation))[:,0]
        
        num_false_head=head_rank.size()[0]#比测试元组分数小的头元组个数
        num_false_tail=tail_rank.size()[0]#比测试元组分数小的尾元组个数
        #num_false_relation=relation_rank.size()[0]#比测试元组分数小的关系元组个数
        
        mean_rank_head=num_false_head
        mean_rank_tail=num_false_tail
        #mean_rank_relation=num_false_relation+1
        
        #计算hit10指标
        if (mean_rank_head<=10):
            hit10_head=1
        if (mean_rank_tail<=10):
            hit10_tail=1
            '''
        if (mean_rank_relation<10):
            hit10_relation=1
            '''
        return mean_rank_head,mean_rank_tail,hit10_head,hit10_tail
    
    
    def test_transE(self):
        hit10_enlity=0
        mean_rank_enlity=0
        
        for test_pair in self.test_tuple:
            mrh,mrt,hth,htt=self.pair_rank(test_pair)
            hit10_enlity+=(hth+htt)/2
            mean_rank_enlity+=(mrh+mrt)/2

            
        hit10_enlity=hit10_enlity/len(self.test_tuple)
        
        mean_rank_enlity=mean_rank_enlity/len(self.test_tuple)
        
        return hit10_enlity,mean_rank_enlity
    
    def test_transH_or_R(self):
        hit10_enlity=0
        mean_rank_enlity=0
        
        for test_pair in self.test_tuple:
            mrh,mrt,hth,htt=self.pair_rank(test_pair)
            hit10_enlity+=(hth+htt)/2
            mean_rank_enlity+=(mrh+mrt)/2

            
        hit10_enlity=hit10_enlity/len(self.test_tuple)
        
        mean_rank_enlity=mean_rank_enlity/len(self.test_tuple)
        
        return hit10_enlity,mean_rank_enlity
 
class Hingleloss(nn.Module):
    def __init__(self):
        super(Hingleloss,self).__init__()
        
    def forward(self,correct_score,corrupt_score,margin):
        #print(correct_score.size())
        zero_tensor = torch.cuda.FloatTensor(correct_score.size())
        zero_tensor.zero_()
        zero_tensor = autograd.Variable(zero_tensor)
        
        #r+correct(h+r-t)-corrupt(h+r-t)
        return torch.sum(torch.max(correct_score-corrupt_score+margin,zero_tensor))
    
    
    def normloss(embedding,dim=1):
        norm = torch.sum(embedding ** 2, dim=dim, keepdim=True)
        return torch.sum(torch.max(norm - autograd.Variable(torch.cuda.FloatTensor([1.0])), autograd.Variable(torch.cuda.FloatTensor([0.0]))))       
        
if __name__=='__main__' :   
    
    
    test_enlity,test_relation,test_pairs=dataloader.get_pairs(dataloader.FB15K_BASE_PATH_test_PATH)
    test_data_load= dataloader.DataLoad(test_enlity,test_relation)
    test_pairs=dataloader.load_without_dic(test_pairs)
    #model=TransE(len(entity),len(relationShips),1,200).cuda()
    t=os.getcwd()+'\model_saved\savedmodel.pkl'
    saved_model=torch.load(t)
    test_part=Evaluation.Evaluation(saved_model,test_pairs)
    hit10_enlity,mean_rank_enlity=test_part.test()
    print('Hit10: {} Mean_rank: {}'.format(hit10_enlity,mean_rank_enlity))
            
        
        
        
        